import os
import pandas as pd
import gradio as gr

import plotly.graph_objects as go

from pathlib import Path
from copy import deepcopy
from country_color import country_color_series


current_path = os.path.abspath(__file__)
current_dir = Path(current_path).parent
parent_dir = Path(current_path).parent.parent

oecd_dir_path = parent_dir / 'oecd_io_data' / 'NATIOTTL' 

current_dir = current_dir / 'data'
csv_names = [oecd_dir_path / s for s in os.listdir(oecd_dir_path)]
csv_names_dic = {s.name:s for s in csv_names}
print(current_dir)

# def read_summary_data():
read_me_path = current_dir / 'OECD_README.xlsx'
data = pd.read_excel(read_me_path, index_col=0)
data = data[data['fit'] == 1]
row_col_data = pd.read_excel(read_me_path, sheet_name='row_col_name')
row_col_to_cn = dict(zip(row_col_data['name'], row_col_data['cn_name']))
cn_to_code = dict(zip(data['国家'],data['code']))
cn_to_en = dict(zip(data['国家'],data['country']))
en_to_cn = dict(zip(data['country'],data['国家']))
cn_to_fit = dict(zip(data['国家'],data['fit']))
code_to_cn = dict(zip(data['code'],data['国家']))


def look_csv_by_country(country):
    country_code = cn_to_code[country]
    csv_country = [s for s in csv_names if country_code in s.name]
    return csv_country


def read_csv(pth):
    """
    此函数用于读取指定路径的 CSV 文件，并对数据进行处理和分割。

    参数:
    pth (str): 要读取的 CSV 文件的路径。

    返回:
    tuple: 包含三个 DataFrame 对象，分别为投入产出矩阵、支出法 GDP 数据和生产法 GDP 数据。
    """
    # 读取指定路径的 CSV 文件，并将第一列作为索引
    data = pd.read_csv(pth, index_col=0)
    # 将数据的列名替换为对应的中文名称
    data.columns = [row_col_to_cn[s] for s in data.columns]
    # 将数据的索引替换为对应的中文名称
    data.index = [row_col_to_cn[s] for s in data.index]
    # 提取投入产出矩阵，即数据的前 45 行和前 45 列
    matrix = data.iloc[:45, :45]
    # 提取支出法 GDP 数据，即数据的前 45 行和第 45 列之后的部分
    expense = data.iloc[:45, 45:]
    # 提取生产法 GDP 数据，即数据的第 45 行之后和前 45 列的部分
    vas = data.iloc[45:, :45]
    return matrix, expense, vas

def plot_df(data_df):
    fig = go.Figure()
    country_count = {}
    for col in data_df.columns:
        cn_country = code_to_cn[col[:3].upper()]
        if cn_country not in country_count:
            country_count[cn_country] = [1]
        else:
            v = country_count[cn_country][-1]
            country_count[cn_country].append(v+1)
        v_color = country_color_series[cn_country][country_count[cn_country][-1]]
        fig.add_trace(go.Bar(x=data_df.index.to_list(), y=data_df[col].to_list(), name=col, marker_color=v_color))
    fig.update_layout(
        barmode='group',
        title='对比',
        yaxis_title='单位：亿美元/(%)',
        height=600,
    )
    return fig

m, e, v = read_csv(look_csv_by_country('中国')[0])


class ShowData:
    def __init__(self):
        # 修改此处，使用 tolist() 方法代替 to_list()
        print('初始化--你现在运行的是查看OECD国家的投入产出表数据')
        self.countrys = data['国家'].tolist()
        self.years = list(range(1995,2021))
        self.csv_names = csv_names
        self.csv_names_dic = csv_names_dic
        self.row_col_to_cn = row_col_to_cn
        self.io_dic = {
            '投入产出矩阵,第一象限':(m.index.to_list(),m.columns.to_list()),
            '生产法GDP,第三象限': (v.index.to_list(),v.columns.to_list()),
            '支出法GDP,第二象限':(e.index.to_list(),e.columns.to_list())
            }
        self.io_dic_number = {'投入产出矩阵,第一象限':0, '生产法GDP,第三象限': 2, '支出法GDP,第二象限':1}

    def plot_data(self, csv_choose, row_col_select, row_select, col_select, data_type):
        print(csv_choose)
        print(row_col_select)
        print(row_select)
        print(col_select)
        n = self.io_dic_number[row_col_select]
        li = []
        for csv in csv_choose:
            data_df = read_csv(self.csv_names_dic[csv])[n]
            se = data_df.loc[row_select, col_select]
            if se.shape[1] > se.shape[0]:
                se = se.T.iloc[:,0]
            print(csv, se.shape)
            se.name = csv
            li.append(se)
        data_choose = pd.concat(li, axis=1)
        data_choose.columns = [s.replace('ttl.csv','') for s in csv_choose]
        data_choose = data_choose / 100 # 百万美元转换为亿美元
        if data_type == '比率':
            data_choose = data_choose / data_choose.sum()
            data_choose = data_choose.round(3)
        else:
            data_choose = data_choose.round(1)
        ot_data = deepcopy(data_choose)
        ot_data.insert(0,'index',ot_data.index.to_list())
        text = f"你现在看的是投入产出表的{n}象限，单位是{data_type}，\n数据来源是{csv_choose},\n行是{row_select},\n列是{col_select}"
        return plot_df(data_choose), ot_data, text

    def select_all_row(self, io_data):
        row, col = self.io_dic[io_data]
        return gr.CheckboxGroup(choices=row,value=row,label='行')

    def deselect_all_row(self, io_data):
        row,col = self.io_dic[io_data]
        return gr.CheckboxGroup(choices=row,value=[],label='行')

    def select_all_col(self, io_data):
        row, col = self.io_dic[io_data]
        return gr.CheckboxGroup(choices=col,value=col,label='列')

    def deselect_all_col(self, io_data):
        row,col = self.io_dic[io_data]
        return gr.CheckboxGroup(choices=col,value=[],label='列')

    def select_io_data(self, item):
        row, col = self.io_dic[item]
        return gr.CheckboxGroup(choices=row, label='行'), gr.CheckboxGroup(choices=col, label='列')

    def search_country(self, country, year):
        if len(country) == 0:
            return gr.CheckboxGroup(choices=[], interactive=True)
        if len(year) == 0:
            return gr.CheckboxGroup(choices=[], interactive=True)
        li = []
        print(country)
        for c in country:
            country_code = cn_to_code[c]
            for doc_name in csv_names:
                if country_code in doc_name.name:
                    for y in year:
                        if str(y) in doc_name.name:
                            li.append(doc_name.name)

        return gr.CheckboxGroup(choices=li, value=li, interactive=True)

    def html(self):
        with gr.Blocks() as demo:
            gr.Markdown("# 多国投入产出对比图")
            with gr.Tab("选择数据"):
                gr.Markdown("## 1.选择国家和年份（可多选）")
                with gr.Row():
                    countrys = gr.CheckboxGroup(label='国家选择', choices=self.countrys, value='中国')
                with gr.Row():
                    years = gr.CheckboxGroup(label='年份选择', choices=self.years, value=2020)
                gr.Markdown("## 2.选择数据文件csv（可多选）")
                with gr.Row():
                    choose_csv = gr.CheckboxGroup(label='选择文件', choices=[], interactive=True)
                with gr.Row():
                    col_row = gr.Radio(
                        label='矩阵', 
                        choices=list(self.io_dic.keys()), 
                        value=list(self.io_dic.keys())[0], interactive=True, scale=1)
                gr.Markdown("## 3.选择行和列（可多选）")
                gr.Markdown("### 注意：如果行多选，列需要单选，如果列多选，行需要单选，总值，不可以两个都多选，否则会报错")
                with gr.Row():
                    row = gr.CheckboxGroup(label='行', choices=m.index.to_list(), interactive=True, scale=5)
                    row_select_all = gr.Button('全选', scale=1)
                    row_select_all.click(self.select_all_row, inputs=col_row, outputs=row)
                    row_deselect_all = gr.Button('全不选', scale=1)
                    row_deselect_all.click(self.deselect_all_row, inputs=col_row, outputs=row)
                with gr.Row():
                    col = gr.CheckboxGroup(label='列', choices=m.index.to_list(), interactive=True, scale=5)
                    col_select_all = gr.Button('全选', scale=1)
                    col_select_all.click(self.select_all_col, inputs=col_row, outputs=col)
                    col_deselect_all = gr.Button('全不选', scale=1)
                    col_deselect_all.click(self.deselect_all_col, inputs=col_row, outputs=col)
                gr.Markdown("## 4.选择数据类型")
                gr.Markdown("### 注意：绝对值单位为亿美元，原始数据为百万美元，比率为每个子部门除以总体")
                with gr.Row():
                    data_type = gr.Radio(label='数据类型', choices=['绝对值', '比率'], value='绝对值', interactive=True, scale=1)
                    plot_bt = gr.Button('生成图表', scale=1)

            with gr.Tab("生成图表"):
                gr.Markdown("## 5.生成图表")
                show_text = gr.Textbox(label='图表信息', interactive=False, lines=5)
                plot_button = gr.Button('生成图表', scale=1)
                plot = gr.Plot(label='多国多年投入产出对比图:')

            with gr.Tab("数据"):
                gr.Markdown("## 6.数据")
                dfdf = gr.DataFrame(label='数据')

            gr.on(
                triggers=[countrys.change, years.change], 
                fn=self.search_country, 
                inputs=[countrys, years], 
                outputs=[choose_csv]
                )
            gr.on(
                triggers=[col_row.change],
                fn=self.select_io_data,
                inputs=[col_row],
                outputs=[row, col]
                )
            gr.on(
                triggers=[plot_bt.click, plot_button.click],
                fn=self.plot_data,
                inputs=[choose_csv, col_row, row, col, data_type],
                outputs=[plot,dfdf,show_text]
                )

        return demo

if __name__ == '__main__':
    sd = ShowData()
    sd.html().launch()